"""
    create by IAmFiveHigh on 2024/8/14
"""


class SegmentTree:
    def __init__(self, arr, merge):
        self._data = arr
        self._size = len(arr)
        self.merge = merge
        # 基础的线段树大小大约为数组大小的4倍
        self._tree = [None] * self._size * 4
        self._build_segment_tree()

    # 创建线段树
    def _build_segment_tree(self):
        self._build_segment_tree_by_index(0, 0, self._size - 1)

    def _build_segment_tree_by_index(self, index: int, l: int, r: int):
        if l >= r:
            # 递归到底的条件这时候是叶结点
            self._tree[index] = self._data[l]
            return

        left_child_index = SegmentTree._left_child_index(index)
        right_child_index = SegmentTree._right_child_index(index)

        # 从中间分开，分别去创建左孩子和右孩子
        mid = int((r - l) / 2) + l
        self._build_segment_tree_by_index(left_child_index, l, mid)
        self._build_segment_tree_by_index(right_child_index, mid + 1, r)
        # 当前index是左孩子和右孩子的merge
        self._tree[index] = self.merge(self._tree[left_child_index], self._tree[right_child_index])
        return

    def get_size(self) -> int:
        return self._size

    def get(self, index: int):
        if index < 0 or index >= self._size:
            raise Exception("Error: index 不合法")
        return self._data[index]

    @staticmethod
    def _left_child_index(index: int) -> int:
        return index * 2 + 1

    @staticmethod
    def _right_child_index(index: int) -> int:
        return index * 2 + 2

    def query(self, query_l: int, query_r: int):
        if query_l < 0 or query_l > query_r or query_r >= self._size:
            raise Exception("Error: 查询的边界不合规")
        return self._query_by_index(0, 0, self._size - 1, query_l, query_r)

    def _query_by_index(self, tree_index: int, l: int, r: int, query_l: int, query_r: int):
        # 如果查询的l等于l 查询的r等于r, 那直接返回当前index的值
        if query_l == l and query_r == r:
            return self._tree[tree_index]

        # 没有正巧查到存好的值
        # 需要的值可能在左孩子里也可能在右孩子里，也可能两边都有一部分
        mid = int((r - l) / 2) + l

        left_child_index = SegmentTree._left_child_index(tree_index)
        right_child_index = SegmentTree._right_child_index(tree_index)

        if query_l >= mid + 1:
            # 查询的左边界都大于中间值了说明查询的值在右孩子中
            return self._query_by_index(right_child_index, mid + 1, r, query_l, query_r)
        elif query_r <= mid:
            # 查询的右边界都小于等于中间值了说明查询的值在左孩子中
            return self._query_by_index(left_child_index, l, mid, query_l, query_r)
        # 到这里的说明查询值部分在左孩子，部分在右孩子
        # 查询左孩子时，query_r要等于mid
        left_result = self._query_by_index(left_child_index, l, mid, query_l, mid)
        # 查询右孩子时，query_l要等于mid + 1
        right_result = self._query_by_index(right_child_index, mid + 1, r, mid + 1, query_r)
        # 返回左边查询到的值和右边查询的值的merge
        return self.merge(left_result, right_result)

    #  更新某个index的元素
    def set(self, index, element):
        if index < 0 or index >= self._size:
            raise Exception("Error: 传入的index不合理")
        self._data[index] = element
        self._set(0, 0, self._size - 1, index, element)

    def _set(self, tree_index: int, l: int, r: int, index: int, element):
        # 找到对应index的时候
        if l == r:
            self._tree[tree_index] = element
            return

        left_child_index = self._left_child_index(tree_index)
        right_child_index = self._right_child_index(tree_index)

        mid = int((r - l) / 2) + l
        # 如果index在左半部分，就去左半部分找
        if index <= mid:
            self._set(left_child_index, l, mid, index, element)
        # 如果index在右半部分，就去右半部分找
        elif index > mid:
            self._set(right_child_index, mid + 1, r, index, element)
        self._tree[tree_index] = self.merge(self._tree[left_child_index], self._tree[right_child_index])

    def __str__(self):

        return str(self._tree)


if __name__ == '__main__':
    def sum1(x, y):
        return x + y


    demo_list = [1, 2, 3, 4, 5, -1, -2, -3, -4, -5]
    segment_tree = SegmentTree(demo_list, sum1)
    print(segment_tree)

    # 把第五位的-1改成0
    segment_tree.set(5, 0)
    print(segment_tree)


